import easyocr
import re


def ocr(file_name):
    # 初始化 EasyOCR，选择语言（中文）和 GPU（如果有的话）
    reader = easyocr.Reader(['ch_sim', 'en'], gpu=True)  # gpu=True 如果您的计算机上有 NVIDIA GPU

    # 加载图片（这里假设图片文件名为 '49-1.jpg'，并且与脚本在同一目录下）
    raw_text = reader.readtext(file_name)

    # 提取股票代码（这里假设股票代码以 '.SZ' 或 '.SH' 结尾，或者为纯数字）
    stock_codes = []
    pattern = re.compile(r'\b\d{6}\b')

    for line in raw_text:
        # line 是一个元组，包含 (bbox, text, confidence)
        text = line[1]
        matches = pattern.findall(text)
        stock_codes.extend(matches)  # 将所有匹配项添加到列表中

    # 返回图片中的所有股票代码
    return stock_codes

